import openai
import json
import os
import time
import numpy as np
import cv2
import re
from typing import Dict, List, Any, Optional

# 配置DeepSeek API
client = openai.OpenAI(
    api_key="****************",  # 替换为你的实际API密钥
    base_url="https://api.deepseek.com/v1"
)


class ImageFusionAgent:
    def __init__(self, fusion_function, stitch_function):
        """
    基于DeepSeek的AI Agent，支持多轮对话和图像融合

    Args:
        fusion_function: 图像融合函数 fuse_color_images(near_img, far_img)
    """
        self.fusion_function = fusion_function
        self.stitch_function = stitch_function

        self.conversation_history = []

        # 系统提示词，定义Agent的角色和能力，包括工具使用说明
        self.system_message = {
            "role": "system",
            "content": """
你是一个专业的图像处理助手，你会根据用户的需求对图像进行不同的处理。

# 工具使用
你可以使用以下工具：
- fuse_images: 当用户提供了两张图像路径并请求融合时使用
- stitch_images: 当用户提供了多张图像路径并请求拼接（配准）时使用

# 响应格式
你的所有响应必须是JSON格式，包含以下字段：
{
    "response": "给用户的自然语言回复",
    "tool": "要使用的工具名称，如'fuse_images'、'stitch_images'，如果没有工具使用则为'none'",
    "tool_input": {
        "images": ["图像路径1", "图像路径2", ...]  # 对于stitch_images，可以有多张图像
    }
}

# 图像路径管理
你需要从对话中提取用户提到的图像路径，并维护在tool_input中：
1. 每次对话都要及时根据上下文正确理解用户意图，更新tool_input中images字段的图像路径
2. 在使用'fuse_images'工具时，如果用户没有明确指定哪张是近景/远景，请根据上下文推断，若无法推断，默认第一张为近景
3. 对于'stitch_images'工具，images字段应该包含所有需要拼接的图像路径列表
4. 如果无法推断，请询问用户澄清
5. 保持路径的准确性，直到工具执行成功

# 注意事项
1. 根据用户意图选择合适的工具：
   - 当用户提到"融合"、"混合"时使用fuse_images，且只支持两张图像
   - 当用户提到"拼接"、"配准"、"对齐"、"接片"时使用stitch_images，支持多张图像
2. 只有在用户明确提供了图像路径并请求处理时，才使用相应工具；如果用户只提供了一张图像或没有提供图像，请继续询问缺失的信息
3. 当需要使用'fuse_images'工具时，不要做出绝对肯定的承诺，告知用户你将进行尝试
4. 在response中不要提到你所收到的路径, 也不要说你将结果保存到了本地
5. 保持对话友好、专业且有帮助
"""
        }
        #

        # 初始化对话历史
        self.conversation_history.append(self.system_message)
        # 存储图像路径（支持多个图像）
        self.image_paths = []  # 改为列表存储多个图像路径

        # 存储当前会话的图像路径
        self.near_image_path = None
        self.far_image_path = None

        self.last_fusion_result = None
        self.last_stitch_result = None

    def chat(self, user_input: str) -> Dict[str, Any]:
        """
    处理用户输入，支持多轮对话

    Args:
        user_input: 用户文本输入

    Returns:
        包含响应和状态的字典
    """
        # 从用户输入中提取可能的图像路径
        # self._extract_image_paths(user_input)

        # 添加用户消息到对话历史
        user_message = {"role": "user", "content": user_input}
        self.conversation_history.append(user_message)

        # 调用DeepSeek API，传递完整的对话历史
        llm_response = self._call_deepseek()

        # 解析LLM的JSON响应
        parsed_response = self._parse_llm_response(llm_response)
        print(parsed_response)

        # 如果LLM决定使用工具，执行相应的工具
        # if parsed_response.get("tool") and parsed_response["tool"] != "none":
        #     tool_result = self._execute_tool(parsed_response)
        #
        #     # 将工具执行结果添加到对话历史
        #     if tool_result:
        #         tool_message = {
        #             "role": "system",
        #             "content": f"工具执行结果: {tool_result}"
        #         }
        #         self.conversation_history.append(tool_message)
        #
        #         # 再次调用LLM，让它基于工具执行结果生成响应
        #         follow_up_response = self._call_deepseek()
        #         parsed_follow_up = self._parse_llm_response(follow_up_response)
        #
        #         # 添加助手响应到对话历史
        #         assistant_message = {"role": "assistant", "content": follow_up_response}
        #         self.conversation_history.append(assistant_message)
        #
        #         return {
        #             "response": parsed_follow_up.get("response", "处理完成"),
        #             "tool_used": parsed_response["tool"],
        #             "tool_result": tool_result,
        #             "tool_input": parsed_response.get("tool_input", {}),
        #             "conversation_id": id(self)
        #         }
        #
        # # 添加助手响应到对话历史
        # assistant_message = {"role": "assistant", "content": llm_response}
        # self.conversation_history.append(assistant_message)
        #
        # # 返回响应
        # return {
        #     "response": parsed_response.get("response", "我没有理解您的请求"),
        #     "tool_used": parsed_response.get("tool", "none"),
        #     "tool_input": parsed_response.get("tool_input", {}),
        #     "conversation_id": id(self)
        # }
        # 添加助手响应到对话历史
        assistant_message = {"role": "assistant", "content": llm_response}
        self.conversation_history.append(assistant_message)

        # 返回初步响应，但标记需要后续处理
        response_data = {
            "response": parsed_response.get("response", "我没有理解您的请求"),
            "tool": parsed_response.get("tool", "none"),
            "tool_input": parsed_response.get("tool_input", {}),
            "conversation_id": id(self),
            "needs_processing": False  # 默认不需要后续处理
        }

        # 如果LLM决定使用工具，标记需要异步处理
        if parsed_response.get("tool") and parsed_response["tool"] != "none":
            response_data["needs_processing"] = True
            # 在响应中添加提示信息
            if "正在处理" not in response_data["response"]:
                response_data["response"] += "\n\n正在处理您的图像，这可能需要一些时间..."

        return response_data
        # return parsed_response

    def _call_deepseek(self) -> str:
        """调用DeepSeek API进行对话，传递完整的对话历史"""
        try:
            response = client.chat.completions.create(
                model="deepseek-chat",
                messages=self.conversation_history,
                temperature=0.7,
                max_tokens=1024,
                stream=False
            )
            return response.choices[0].message.content
        except Exception as e:
            return json.dumps({
                "response": f"抱歉，我遇到了一些问题：{str(e)}。请稍后重试。",
                "tool": "none",
                "tool_input": {}
            })

    def _parse_llm_response(self, llm_response: str) -> Dict[str, Any]:
        """解析LLM的JSON响应"""
        try:
            # 尝试直接解析JSON
            return json.loads(llm_response)
        except json.JSONDecodeError:
            # 如果直接解析失败，尝试从文本中提取JSON
            json_match = re.search(r'\{.*\}', llm_response, re.DOTALL)
            if json_match:
                try:
                    return json.loads(json_match.group())
                except json.JSONDecodeError:
                    pass

            # 如果无法提取JSON，返回默认响应
            return {
                "response": llm_response,
                "tool": "none",
                "tool_input": {}
            }

    # def _extract_image_paths(self, user_input: str):
    #     """从用户输入中提取图像路径（支持多个）"""
    #     # 简单的路径提取逻辑 - 可以根据需要扩展
    #     paths = re.findall(r'([a-zA-Z]:\\[^\s]+|/(?:[^/\s]+/)*[^/\s]+|\./[^\s]+)', user_input)
    #
    #     for path in paths:
    #         if os.path.exists(path) and path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
    #             # 添加到图像路径列表（避免重复）
    #             if path not in self.image_paths:
    #                 self.image_paths.append(path)

    def _execute_tool(self, parsed_response: Dict[str, Any]) -> str:
        """执行工具函数"""
        tool = parsed_response.get("tool")
        tool_input = parsed_response.get("tool_input", {})

        if tool == "fuse_images":
            images = tool_input.get("images", [])
            if len(images) != 2:
                return "错误：图像融合需要恰好两张图像"

            image1, image2 = images[0], images[1]

            # 验证图像路径
            for img_path in [image1, image2]:
                if not os.path.exists(img_path):
                    return f"错误：图像路径不存在 - {img_path}"

            try:
                # 调用融合函数
                success, fused_image = self.fusion_function(image1, image2)

                # 保存融合结果
                self.last_fusion_result = (success, fused_image)

                if success:
                    # 清空已处理的图像路径
                    self.image_paths = []
                    return "图像融合成功完成"
                else:
                    return "图像融合失败，请检查图像文件格式和内容"

            except Exception as e:
                return f"图像融合过程中出现错误：{str(e)}"

        elif tool == "stitch_images":
            # 拼接处理（支持多张图像）
            images = tool_input.get("images", [])

            if len(images) < 2:
                return "错误：图像拼接至少需要两张图像"

            # 验证所有图像路径
            for img_path in images:
                if not os.path.exists(img_path):
                    return f"错误：图像路径不存在 - {img_path}"

            try:
                # 调用拼接函数（现在应该接受图像路径列表）
                success, stitched_image = self.stitch_function(images)

                # 保存拼接结果
                self.last_stitch_result = (success, stitched_image)

                if success:
                    # 清空已处理的图像路径
                    self.image_paths = []
                    return f"图像拼接成功完成，共处理了{len(images)}张图像"
                else:
                    return "图像拼接失败，请检查图像文件格式和内容"

            except Exception as e:
                return f"图像拼接过程中出现错误：{str(e)}"

        return f"未知工具：{tool}"

    # def get_conversation_summary(self) -> Dict[str, Any]:
    #     """获取当前对话摘要"""
    #     return {
    #         "near_image": self.near_image_path,
    #         "far_image": self.far_image_path,
    #         "conversation_turns": len([msg for msg in self.conversation_history if msg["role"] == "user"]),
    #         "has_both_images": bool(self.near_image_path and self.far_image_path)
    #     }
    def get_conversation_summary(self) -> Dict[str, Any]:
        """获取当前对话摘要"""
        # 从最新的LLM响应中获取当前图像路径
        latest_response = None
        for msg in reversed(self.conversation_history):
            if msg["role"] == "assistant":
                try:
                    latest_response = json.loads(msg["content"])
                    break
                except:
                    continue

        near_img = None
        far_img = None
        if latest_response and "tool_input" in latest_response:
            near_img = latest_response["tool_input"].get("near_img")
            far_img = latest_response["tool_input"].get("far_img")

        return {
            "near_image": near_img,
            "far_image": far_img,
            "conversation_turns": len([msg for msg in self.conversation_history if msg["role"] == "user"]),
            "has_both_images": bool(near_img and far_img)
        }

    def clear_conversation(self):
        """清空对话历史（保留系统消息）"""
        self.conversation_history = [self.system_message]
        self.near_image_path = None
        self.far_image_path = None


def create_gaussian_filter(shape, cutoff_ratio=0.1):
    """
    创建高斯低通滤波器

    参数:
        shape: 滤波器尺寸 (rows, cols)
        cutoff_ratio: 截止频率比例，控制低频区域大小
    """
    rows, cols = shape
    crow, ccol = rows // 2, cols // 2

    # 创建距离矩阵
    u = np.arange(rows)
    v = np.arange(cols)
    u, v = np.meshgrid(u, v, indexing='ij')

    # 计算每个点到中心的距离
    d = np.sqrt((u - crow) ** 2 + (v - ccol) ** 2)

    # 创建高斯低通滤波器
    cutoff = min(rows, cols) * cutoff_ratio
    gaussian_lowpass = np.exp(-(d ** 2) / (2 * (cutoff ** 2)))

    return gaussian_lowpass


def fuse_color_images(near_img_path, far_img_path):
    """
    融合两幅彩色图像

    参数:
        near_img: 近景图像 (彩色)
        far_img: 远景图像 (彩色)
        cutoff_ratio: 截止频率比例

    返回:
        fused_img: 融合后的图像
        near_high_freq: 近景图像高频信息 (用于可视化)
        far_low_freq: 远景图像低频信息 (用于可视化)
    """
    cutoff_ratio = 0.06
    try:
        near_img = cv2.imread(near_img_path)  # 替换为你的近景图像路径
        far_img = cv2.imread(far_img_path)  # 替换为你的远景图像路径
        # 确保图像尺寸相同
        if near_img.shape != far_img.shape:
            far_img = cv2.resize(far_img, (near_img.shape[1], near_img.shape[0]))

        # 获取图像尺寸
        rows, cols, channels = near_img.shape

        # 创建高斯滤波器
        gaussian_filter = create_gaussian_filter((rows, cols), cutoff_ratio)

        # 初始化结果数组
        fused_channels = []
        near_high_freq_channels = []
        far_low_freq_channels = []

        # 对每个颜色通道分别处理
        for ch in range(channels):
            # 傅里叶变换并中心化
            near_fft = np.fft.fft2(near_img[:, :, ch].astype(np.float32))
            far_fft = np.fft.fft2(far_img[:, :, ch].astype(np.float32))

            near_fft_shift = np.fft.fftshift(near_fft)
            far_fft_shift = np.fft.fftshift(far_fft)

            # 使用高斯滤波器提取低频信息
            near_low_freq = near_fft_shift * gaussian_filter
            far_low_freq = far_fft_shift * gaussian_filter

            # 通过原图像减去低频信息得到高频信息
            near_high_freq = near_fft_shift - near_low_freq

            # 融合：近景的高频 + 远景的低频
            fused_freq = near_high_freq + far_low_freq

            # 逆傅里叶变换
            fused_shift = np.fft.ifftshift(fused_freq)
            fused_ch = np.fft.ifft2(fused_shift)
            fused_ch = np.abs(fused_ch)

            # 收集结果
            fused_channels.append(fused_ch)
            near_high_freq_channels.append(near_high_freq)
            far_low_freq_channels.append(far_low_freq)

        # 合并通道
        fused_img = np.stack(fused_channels, axis=2)
        fused_img = np.uint8(np.clip(fused_img, 0, 255))

        # 将频域信息转换为可视化格式（仅使用一个通道用于显示）
        # near_high_freq_vis = np.log(1 + np.abs(near_high_freq_channels[0]))
        # far_low_freq_vis = np.log(1 + np.abs(far_low_freq_channels[0]))

        # cv2.imwrite('fused_img.jpg', fused_img)
        return True, fused_img

    except Exception as e:
        return False, None
# 示例使用代码
if __name__ == "__main__":

    # 创建Agent实例
    agent = ImageFusionAgent(fuse_color_images)

    print("图像融合AI助手已启动！输入'退出'结束对话，输入'清空'重置对话")

    while True:
        user_input = input("\n您: ").strip()

        if user_input.lower() in ['退出', 'exit', 'quit']:
            break
        elif user_input.lower() in ['清空', 'clear', '重置']:
            agent.clear_conversation()
            print("对话已清空")
            continue

        # 处理用户输入
        response = agent.chat(user_input)

        print(f"\n助手: {response['response']}")

        if response.get('tool_used') and response['tool_used'] != 'none':
            print(f"[工具使用] 使用了 {response['tool_used']} 工具")
            print(f"[工具结果] {response.get('tool_result', '')}")

        # 显示当前状态
        # summary = agent.get_conversation_summary()
        # if summary['near_image'] or summary['far_image']:
        #     images = []
        #     if summary['near_image']:
        #         images.append(f"近景: {summary['near_image']}")
        #     if summary['far_image']:
        #         images.append(f"远景: {summary['far_image']}")
        #     print(f"[状态] 已收集: {', '.join(images)}")
        summary = agent.get_conversation_summary()
        if summary['near_image'] or summary['far_image']:
            images = []
            if summary['near_image']:
                images.append(f"近景: {summary['near_image']}")
            if summary['far_image']:
                images.append(f"远景: {summary['far_image']}")
            print(f"[状态] 已收集: {', '.join(images)}")